import pyaudio
import wave
import time
import sys


def test_microphone(device_index, record_seconds=5):
    CHUNK = 1024
    FORMAT = pyaudio.paInt16
    CHANNELS = 2
    RATE = 44100

    p = pyaudio.PyAudio()

    try:
        # 检查设备是否可用
        device_info = p.get_device_info_by_index(device_index)
        if device_info['maxInputChannels'] < 1:
            print(f"设备 {device_index} 不是输入设备或不可用")
            return

        print(f"\n准备测试设备 {device_index} ({device_info['name']})...")
        print("将在3秒后开始录音，请准备...")

        # 3秒倒计时
        for i in range(3, 0, -1):
            print(f"{i}...", end=' ')
            sys.stdout.flush()
            time.sleep(1)
        print("开始录音!")

        stream = p.open(
            input_device_index=device_index,
            format=FORMAT,
            channels=CHANNELS,
            rate=RATE,
            input=True,
            frames_per_buffer=CHUNK
        )

        frames = []
        print(f"正在录音 {record_seconds} 秒... (按Ctrl+C可提前结束)")

        try:
            start_time = time.time()
            while time.time() - start_time < record_seconds:
                data = stream.read(CHUNK, exception_on_overflow=False)
                frames.append(data)
                # 显示进度条
                elapsed = time.time() - start_time
                progress = int(elapsed / record_seconds * 20)
                print(f"\r[{'#' * progress}{'-' * (20 - progress)}] {elapsed:.1f}/{record_seconds}s", end='')

            print("\n录音完成!")

        except KeyboardInterrupt:
            print("\n录音被用户中断")
        except Exception as e:
            print(f"\n录音过程中发生错误: {str(e)}")
        finally:
            stream.stop_stream()
            stream.close()

        # 保存文件
        if frames:
            filename = f"test_{device_index}_{int(time.time())}.wav"
            with wave.open(filename, 'wb') as wf:
                wf.setnchannels(CHANNELS)
                wf.setsampwidth(p.get_sample_size(FORMAT))
                wf.setframerate(RATE)
                wf.writeframes(b''.join(frames))
            print(f"已保存测试录音: {filename}")

    except Exception as e:
        print(f"初始化录音设备时出错: {str(e)}")
    finally:
        p.terminate()


# 使用方法
if __name__ == "__main__":
    # 列出所有音频设备
    p = pyaudio.PyAudio()
    print("可用的音频输入设备:")
    for i in range(p.get_device_count()):
        dev = p.get_device_info_by_index(i)
        if dev['maxInputChannels'] > 0:
            print(f"{i}: {dev['name']} (输入通道: {dev['maxInputChannels']})")
    p.terminate()

    # 让用户选择设备
    try:
        device_index = int(input("\n请输入要测试的设备索引号: "))
        record_time = int(input("请输入录音时长(秒，默认5秒): ") or 5)
        test_microphone(device_index, record_time)
    except ValueError:
        print("请输入有效的数字")